/*******************************************************************************
 * Copyright 2010 DemandSpot
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ******************************************************************************/
package com.demandspot.aws.sqs;

import java.io.BufferedWriter;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStreamWriter;
import java.io.Serializable;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.RejectedExecutionException;
import java.util.zip.Deflater;
import java.util.zip.DeflaterOutputStream;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;
import java.util.zip.InflaterInputStream;

import org.apache.commons.codec.binary.Base64;
import org.apache.log4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import com.demandspot.aws.AWSConsts;
import com.xerox.amazonws.sqs2.Message;
import com.xerox.amazonws.sqs2.MessageQueue;
import com.xerox.amazonws.sqs2.QueueService;
import com.xerox.amazonws.sqs2.SQSException;
import java.io.BufferedReader;
import java.io.InputStreamReader;

public class AWSSqsService {
	private static final Logger logger = Logger.getLogger(AWSSqsService.class);
	
	@Autowired private ThreadPoolTaskExecutor taskExecutor;
	private List<SqsMessagePoller> messagePollers;
	private Map<SqsMessageHandler, SqsMessagePoller> handlersMap;
	
	private QueueService queueService; 
	
	public AWSSqsService() throws FileNotFoundException, IOException {		
		this(AWSSqsService.class.getResourceAsStream("/conf/aws.properties"));
	}

	public AWSSqsService(InputStream propInputStream) throws IOException {
		Properties props = new Properties();
		props.load(propInputStream);
		init(props);
	}
	
	public AWSSqsService(File propsFile) throws FileNotFoundException, IOException {
		Properties props = new Properties();
		props.load(new FileInputStream(propsFile));
		init(props);
	}
	
	public AWSSqsService(URL propsURL) throws FileNotFoundException, IOException {
		this(new File(propsURL.getFile()));
	}
	
	public AWSSqsService(Properties props) {
		init(props);
	}

	public AWSSqsService(String accessKey, String secretKey) {
		Properties props = new Properties();
		props.put(AWSConsts.AWS_ACCESS_KEY, accessKey);
		props.put(AWSConsts.AWS_SECRET_KEY, secretKey);
		init(props);
	}
	
	public AWSSqsService(String accessKey, String secretKey, boolean secure) {
		Properties props = new Properties();
		props.put(AWSConsts.AWS_ACCESS_KEY, accessKey);
		props.put(AWSConsts.AWS_SECRET_KEY, secretKey);
		props.put(AWSConsts.AWS_SECURE, String.valueOf(secure));
		init(props);
	}
	
	private void init(Properties props) {
		String accessKey = props.getProperty(AWSConsts.AWS_ACCESS_KEY);
		String secretKey = props.getProperty(AWSConsts.AWS_SECRET_KEY);
		String isSecureStr = props.getProperty(AWSConsts.AWS_SECURE, "false");
		boolean isSecure = Boolean.parseBoolean(isSecureStr);
		
		accessKey = System.getProperty(AWSConsts.AWS_ACCESS_KEY, accessKey);
		secretKey = System.getProperty(AWSConsts.AWS_SECRET_KEY, secretKey);
		
		queueService = new QueueService(accessKey, secretKey, isSecure);
		messagePollers = Collections.synchronizedList(new ArrayList<SqsMessagePoller>());
		handlersMap = Collections.synchronizedMap(new HashMap<SqsMessageHandler, SqsMessagePoller>());
	}
	
	public void addMessageHandler(SqsMessageHandler messageHandler, String queueName, int numOfMessages, boolean create, boolean encode) throws SQSException {
		addMessageHandler(messageHandler, queueName, numOfMessages, create, encode, -1);
	}
	
	public void addMessageHandler(SqsMessageHandler messageHandler, String queueName, int numOfMessages, boolean create) throws SQSException {
		addMessageHandler(messageHandler, queueName, numOfMessages, create, true, -1);
	}
	
	public void addMessageHandler(SqsMessageHandler messageHandler, String queueName, int numOfMessages, boolean create, boolean encode, int hideInSecs) throws SQSException {
		MessageQueue messageQueue = null;
		
		if (create) {
			messageQueue = queueService.getOrCreateMessageQueue(queueName);
		} else {
			messageQueue = queueService.getMessageQueue(queueName);
		}
		
		messageQueue.setEncoding(encode);
		messageHandler.setQueue(messageQueue);
		
		SqsMessagePoller poller = new SqsMessagePoller(messageQueue, messageHandler, queueName, numOfMessages, hideInSecs);
		messagePollers.add(poller);
	}
	
	public MessageQueue getMessageQueue(String queueName) throws SQSException {
		return getMessageQueue(queueName, true);
	}
	
	public MessageQueue getMessageQueue(String queueName, boolean create) throws SQSException {
		if (create) {
			return queueService.getOrCreateMessageQueue(queueName);
		} else {
			return queueService.getMessageQueue(queueName);
		}
	}
	
	@Scheduled(fixedRate=60000)
	public void pollMessages() {
		for (SqsMessagePoller poller : messagePollers) {
			try {
				if (!poller.isRunning()) {
					poller.setRunning(true);
					taskExecutor.execute(poller);
				}
			} catch (RejectedExecutionException e) {
				logger.warn("Poller was rejected: " + poller);
			}
		}
	}
	
	public void removeMessageHandler(SqsMessageHandler messageHandler) {
		SqsMessagePoller messagePoller = handlersMap.remove(messageHandler);
		
		if (messagePoller != null) {
			messagePollers.remove(messagePoller);
		}
	}
	
	public static void sendMessage(MessageQueue messageQueue, String content, boolean compress) throws SQSException {
		if (compress) {
			content = compress(content);	
		}
		messageQueue.sendMessage(content);
	}
	
	public static String readCompressedMessage(Message message) throws SQSException {
		try {
			byte[] bytes = Base64.decodeBase64(message.getMessageBody().getBytes());
			BufferedReader reader = new BufferedReader(new InputStreamReader(new InflaterInputStream(new ByteArrayInputStream(bytes))));
			StringBuilder builder = new StringBuilder();
			String line;
			while ((line = reader.readLine()) != null) {
				builder.append(line).append("\n");
			}
			return builder.toString();
		} catch (IOException e) {
			throw new SQSException("failed to read compressed message", e);
		}
	}
	
	public static void sendObjectMessage(MessageQueue messageQueue, Serializable content, boolean compress) throws SQSException {
		byte[] bytes = new byte[0];
		try {
			ByteArrayOutputStream out = new ByteArrayOutputStream();
			ObjectOutputStream oos;
			
			if (compress) {
				GZIPOutputStream gzip = new GZIPOutputStream(out);
				oos = new ObjectOutputStream(gzip);
			} else {
				oos = new ObjectOutputStream(out);
			}
			
			oos.writeObject(content);
			oos.close();
			
			bytes = Base64.encodeBase64(out.toByteArray());
			String str = new String(bytes);
			messageQueue.sendMessage(str);
		} catch (Exception e) {
			throw new SQSException("Failed to send Object (message size " + bytes.length + " bytes)", e);
		}
	}
	
	@SuppressWarnings("unchecked")
	public static <T extends Serializable> T readObjectMessage(Message message, boolean compressed) throws SQSException {
		T toReturn = null;
		try {
			byte[] bytes = Base64.decodeBase64(message.getMessageBody().getBytes());
			ByteArrayInputStream in = new ByteArrayInputStream(bytes);
			ObjectInputStream ins;
			
			if (compressed) {
				GZIPInputStream gzip = new GZIPInputStream(in);
				ins = new ObjectInputStream(gzip);
			} else {
				ins = new ObjectInputStream(in);
			}
			
			toReturn = (T) ins.readObject();
			ins.close();
		} catch (Exception e) {
			throw new SQSException("Failed to read Object", e);
		}
		
		return toReturn;
	}
	
	public static void clearQueue(MessageQueue queue) throws SQSException {
		Message[] messages = new Message[0];
		do {
			messages = queue.receiveMessages(10);
			
			for (Message message : messages) {
				queue.deleteMessage(message);
			}
			
		} while (messages != null && messages.length > 0);
	}
	
	private static String compress(String toCompress) throws SQSException {
		try {
			ByteArrayOutputStream out = new ByteArrayOutputStream();
			Deflater deflater = new Deflater(9);
			DeflaterOutputStream gstream = new DeflaterOutputStream(out, deflater);
			BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(gstream));
			writer.write(toCompress);
			writer.close();
			
			byte[] bytes = Base64.encodeBase64(out.toByteArray());
			toCompress = new String(bytes);
		} catch (IOException e) {
			throw new SQSException("failed to compress content", e);
		}
		return toCompress;
	}

	public List<MessageQueue> listQueues() throws SQSException {
		return queueService.listMessageQueues(null);
	}
}
